from UserDict import *
import copy
from string import Template

#compile_out outputs the appropriate assembly template for each function. 

#exception for when a symbole is used  twice
class SymbolAlreadyUsedError(Exception):
	def __init__(self,symbolname):
		self.symbolname = symbolname
	def __str__(self):
		return repr(self.symbolname)

#exception for a refernce to an undefined symbol
class SymbolNotDefinedError(Exception):
	def __init__(self,symbolname):
		self.symbolname = symbolname
	def __str__(self):
		return repr(self.symbolname)

class stack(list):
	def peek(self):
		return self[-1]
	def push(self,item):
		self.append(item)

#Set for checking uniqueness of symbols in current scope.
class uniquedict(IterableUserDict):
	def __setitem__(self,key,value):
		if key in self or key == "min":
			raise SymbolAlreadyUsedError(key)
		IterableUserDict.__setitem__(self,key,value)
	def __getitem__(self,key):
		if key not in self:
			raise SymbolNotDefinedError(key)
		return IterableUserDict.__getitem__(self,key)

output = """
.text
.align 4,0x90
"""
#Global stack counter
func_count = 0
#Initialise the stack
stk = stack()
#Global loop counter 
loop_count = 0
#Number of local vars
n_lvars = 0


def add_prog(root):
	global func_count,stk,output
	#push a set for checking reuse in the current scope. Will also serve as a store for symbols' addresses. Or
	#at this point (too lazy to pre compute) -  the asm for storing those addresses in eax.
	stk.push(uniquedict())
	#so we need to find constants. will use tupple for later convenience.
	consts = set(find_constants(root))
	#add the address calculation code of these to the global scope (these will be shared between functions, but it doesnt matter)
	for item in consts:
		if str(item[1]) not in stk.peek():
			stk.peek()[str(item[1])] =Template("###location of " + str(item[1]) + "\n	movl	$$.const%s, $destreg"%str(item[1]))
			#should throw exeption
			#print stk.peek()["test"]
			#works!
			#print stk.peek()[str(item[1])]
	for child in root.children:
		funcname = str(child.children[0])
		add_func(child)
		func_count=func_count + 1
	#were now out of scope.
	stk.pop()
	#time for the postambles.
	#only if we actually have consts.
	if len(consts) != 0:
		output+="""
.data
.align 16"""
		for item in consts:
			output +="""
.const%s:"""%str(item[1])
			for lol in "lool":#lol
				output+="""
	.float	%s"""%(str(item[1])+ (".0" if item[0] == "int" else ""))#if its an in append .0, otherwise do nothing

def add_func(func):
	funcname = str(func.children[0])
	stk.peek()[funcname] = "exists"
	global output,n_lvars
	n_lvars = 0
	output+= """
.global %s
%s:
	pushl	%%ebp
	movl	%%esp, %%ebp
	pushl	%%ebx
	"""%(funcname,funcname)
	#work out the number of local vars. Do this search twice cause im lazy
	for child in func.children[1:]:
		child_type = str(child)
		if child_type == "DECL":
			n_lvars = len(child.children)
	#cool. now reserve space for them. Only if we actually need to
	if n_lvars != 0:
		output += """
	#making space for local vars
	movl	8(%%ebp), %%eax
	imull	$4, %%eax, %%eax
	addl	$16, %%eax
	imull	$%i, %%eax, %%eax
	subl	%%eax, %%esp
	andl	$-16, %%esp

	"""%n_lvars
	#we need to start adding stuff from the function body here
	#start from 1 as the functions name is at index 0
	#we have a new scope for the params which is shared with the declarations(no new names are declared in statements)
	#however we still dont need to pop until the end of the function.
	stk.push(copy.deepcopy(stk.peek()))
	for child in func.children[1:]:
		child_type = str(child)
		if child_type == "PARAM":
			add_param(child)
		elif child_type == "DECL":
			add_decl(child)
		elif child_type == "STATEMENTS":
			add_statements(child)
	stk.pop()
	output+="""
	#finished with this function.
	movl	-4(%ebp), %ebx
	leave
	ret
"""
def add_param(param):
	n = 1
	global func_count
	for child in param.children:
		#stk.peek().add(str(child))
		#save the code for getting the address of the nth variable. Put it in edx
		stk.peek()[str(child)] = Template("###location of " + str(child) + "\n	movl	" + str(8+4*n)+ " (%ebp), $destreg")
		n = n + 1
	global output
	#for child in param.children:
	output+= """param%i:
	"""%func_count


def add_decl(decl):
	n = 1
	for child in decl.children:
		#getting the address of the nth local var. too ezmode.
		stk.peek()[str(child)] = Template("###location of localvar " + str(n) + """\n
	movl	8(%%ebp), $destreg
	imull	$$4, $destreg, $destreg
	addl	$$16, $destreg
	imull	$$%i, $destreg, $destreg
	subl	%%ebp, $destreg
	negl	$destreg
	andl	$$-16, $destreg
	"""%n)
		n = n+1
	global output
	output+="""decl%i:
	"""%func_count

def add_statements(statements):
	#lazily checking if there are children, if not return.
	if len(statements.children) == 0:
		return

	global output
	output+="""statements%i:
	"""%func_count
	for child in statements.children:
		if str(child)=="=":
			add_expr(stk.peek()[str(child.children[0])],child.children[1])
		elif str(child.children[0])=="PARAM": #local var
			params = []
			for greatgrandchild in child.children[0].children:
				params.append(str(greatgrandchild))
			func_call(str(child),params)

def add_expr(dest,parent):
	operators = ["+","-","/","*","min"]
	global stk
	if str(parent) == "FLT" or str(parent) == "INT":
		assign(dest,stk.peek()[str(parent.children[0])])
	elif str(parent) in operators:
		#call operator with the factors
		#print str(parent.children[0].children[0])
		#print stk.peek()[str(parent.children[0].children[0])]
		
		if str(parent.children[0]) == "FLT" or str(parent.children[0]) == "INT":
			lhs_str = str(parent.children[0].children[0])
		else:
			lhs_str = str(parent.children[0])
		if str(parent.children[1]) == "FLT" or str(parent.children[1]) == "INT":
			rhs_str = str(parent.children[1].children[0])
		else:
			rhs_str = str(parent.children[1])

		
		
		if lhs_str not in stk.peek():
			add_lhs = new_local_var()
			add_expr(add_lhs,parent.children[0])
		else:
			add_lhs = stk.peek()[lhs_str]
		if rhs_str not in stk.peek():
			add_rhs = new_local_var()
			add_expr(add_rhs,parent.children[1])
		else:
			add_rhs = stk.peek()[rhs_str]
		if True:#not in stk.peek()
			pass#reserve space for a var
			#call the expr with the address of that var as dest
			#call operator(dest,that new address, other child, operator)
		try:
			return operator(dest,add_lhs,add_rhs,str(parent))
		except Exception:
			print add_lhs
	elif str(parent) in stk.peek():
		assign(dest,stk.peek()[str(parent)])
	else:
		#symbols not found anywhere then. raise an exception.
		raise SymbolNotDefinedError(str(parent))


#first arg: code for loading lhs into %s
#second arg: code for loading rhs into %s
def operator(dest,lhs,rhs,operator):
	operator_func = {}
	global output,loop_count
	operator_func["+"] = "addps"
	operator_func["-"] = "subps"
	operator_func["/"] = "divps"
	operator_func["*"] = "mulps"
	operator_func["min"] = "minps"
	output+="""
	#addition
	
	#put the addresses of the source
	#vectors into eax and ebx.
	"""+ rhs.substitute(destreg='%eax') + """
	"""+ lhs.substitute(destreg='%ebx') + """
	
	#put the address of the dest into edx
	"""+dest.substitute(destreg='%edx') + Template("""
	
	#Plus: loop through and copy it all to the variable being held at edx
	movl	8(%%ebp), %%ecx
	shrl	$$2, %%ecx
	jz		.loop_end$lpcnt
	.loop_begin$lpcnt:
		movaps	(%%eax), %%xmm0
		movaps	(%%ebx), %%xmm1
		%s	%%xmm0, %%xmm1
		movaps	%%xmm1, (%%edx)""").substitute(lpcnt = str(loop_count))%operator_func[operator]

	if "cons" not in rhs.substitute(destreg=''):
		output+="""
		addl	$16,	%eax"""
	if "cons" not in lhs.substitute(destreg=''):
		output+="""
		addl	$16,	%ebx"""
	output+=Template("""
		addl	$$16,	%edx
		loopl	.loop_begin$lpcnt
	.loop_end$lpcnt:
	""").substitute(lpcnt = str(loop_count))
	loop_count = loop_count + 1


def find_constants(node):
	if len(node.children) == 0:
		return []
	elif str(node) == "INT":
		return [("int",str(node.children[0]))]
	elif str(node) == "FLT":
		return [("float",str(node.children[0]))]
	else:
		ret = []
		for child in node.children:
			ret = ret + find_constants(child)
		return ret

def assign(dest_add,source_add):
	global output,loop_count
	output+="""
	#assignment
	
	#get source into eax
	""" + source_add.substitute(destreg='%eax') + """
	
	###get dest into edx
	""" + dest_add.substitute(destreg='%edx') + Template("""
	movl	8(%ebp), %ecx
	shrl	$$2,	%ecx
	jz		.loop_end$lpcnt

	.loop_begin$lpcnt:
		movaps	(%eax), %xmm0
		movaps	%xmm0, (%edx)
		""").substitute(lpcnt=str(loop_count))

	if "cons" not in source_add.substitute(destreg=''):
		output+="""addl	$16, %eax"""
	output+=Template("""
		addl	$$16, %edx
		loopl	.loop_begin$lpcnt

	.loop_end$lpcnt:
	""").substitute(lpcnt=str(loop_count))
	loop_count = loop_count + 1

def func_call(funcname,params):
	global output
	output+="""
	######################################
	#got us a func call for """+ funcname + """ we did.
	#params: """ + str(params) + """
	######################################
	#invoking function
	subl	$""" + str(4*(len(params)+1)) + """, %esp
	
	"""
	#hopefully this is the correct range()
	for i in range(1,len(params)+1):
		output+=stk.peek()[params[len(params)-i]].substitute(destreg='%eax')+"""
	movl	%eax, """ + str(4*(len(params)-i+1)) + """(%esp)
	"""

	output+="""#length into this func's stack.
	movl	8(%ebp), %eax
	movl	%eax, 0(%esp)

	call	""" + funcname + """
	#now after doing that chuck the stack back to where it used to be
	addl	$""" + str(4*(len(params)-i+1)) + """, %esp

	"""

def new_local_var():
	global output,n_lvars
	output+="""
	#making space for local vars
	movl	8(%ebp), %eax
	imull	$4, %eax, %eax
	addl	$16, %eax
	#pretty sure this line doesnt matter
	#imull	$1, %eax, %eax
	subl	%eax, %esp
	andl	$-16, %esp


	"""
	n_lvars = n_lvars + 1

	return Template("""
	###location of a tempvar
	movl	8(%%ebp), $destreg
	imull	$$4, $destreg, $destreg
	addl	$$16, $destreg
	imull	$$%i, $destreg, $destreg
	subl	%%ebp, $destreg
	negl	$destreg
	andl	$$-16, $destreg
	"""%(n_lvars+1))


